import asyncio
import logging
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.db import get_async_session, User, Role, UserRole
from app.models.user import get_user_manager, UserManager
from app.utils.auth import get_user_role_codes

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 默认管理员信息
ADMIN_USERNAME = "admin"
ADMIN_EMAIL = "admin@example.com"
ADMIN_PASSWORD = "admin123"  # 在生产环境中应使用环境变量或更安全的方式

async def init_roles(session: AsyncSession):
    """初始化角色"""
    # 检查管理员角色是否存在
    admin_role_query = select(Role).where(Role.name == "admin")
    result = await session.execute(admin_role_query)
    admin_role = result.scalar_one_or_none()
    
    if not admin_role:
        logger.info("创建管理员角色")
        admin_role = Role(
            name="admin",
            display_name="管理员",
            description="系统管理员，拥有所有权限"
        )
        session.add(admin_role)
    
    # 检查普通用户角色是否存在
    user_role_query = select(Role).where(Role.name == "user")
    result = await session.execute(user_role_query)
    user_role = result.scalar_one_or_none()
    
    if not user_role:
        logger.info("创建普通用户角色")
        user_role = Role(
            name="user",
            display_name="普通用户",
            description="普通用户，拥有基本权限"
        )
        session.add(user_role)
    
    await session.commit()
    return admin_role, user_role

async def create_admin_user(user_manager: UserManager, session: AsyncSession):
    """创建管理员用户"""
    # 检查管理员用户是否存在
    admin_query = select(User).where(
        (User.username == ADMIN_USERNAME) | (User.email == ADMIN_EMAIL)
    )
    result = await session.execute(admin_query)
    admin_user = result.scalar_one_or_none()
    
    if not admin_user:
        logger.info(f"创建管理员用户: {ADMIN_USERNAME}")
        admin_user = await user_manager.create(
            {
                "username": ADMIN_USERNAME,
                "email": ADMIN_EMAIL,
                "password": ADMIN_PASSWORD,
                "is_active": True,
                "is_superuser": True,
                "is_verified": True,
                "nickname": "系统管理员"
            }
        )
    
    # 确保管理员角色存在
    admin_role, _ = await init_roles(session)
    
    # 检查用户是否已分配管理员角色
    user_role_query = select(UserRole).where(
        (UserRole.user_id == str(admin_user.id)) & (UserRole.role_id == admin_role.id)
    )
    result = await session.execute(user_role_query)
    user_role = result.scalar_one_or_none()
    
    if not user_role:
        logger.info(f"为用户 {ADMIN_USERNAME} 分配管理员角色")
        user_role = UserRole(
            user_id=str(admin_user.id),
            role_id=admin_role.id
        )
        session.add(user_role)
        await session.commit()
    
    # 验证角色分配是否成功
    roles = await get_user_role_codes(str(admin_user.id), session)
    logger.info(f"用户 {ADMIN_USERNAME} 的角色: {roles}")

async def main():
    """初始化数据的主函数"""
    logger.info("开始初始化数据...")
    
    # 获取会话和用户管理器
    async for session in get_async_session():
        try:
            async for user_manager in get_user_manager(session):
                await create_admin_user(user_manager, session)
                break
        except Exception as e:
            logger.error(f"初始化数据时出错: {str(e)}")
            raise
    
    logger.info("数据初始化完成")

if __name__ == "__main__":
    asyncio.run(main()) 